#!/usr/bin/env python
# coding=utf-8

import re
import argparse
import sys


parser = argparse.ArgumentParser(
    description='''
    有时 pasa出来的结果会出现 移码错误 需要找到原有的gff版本进行校正
    用法:
    pasa_delete_merge.py -i tmp3.gff3 -o Cves_OGS_v3.2.gff3
    ''')




parser.add_argument('-i',
                help='必须给定，输入的gff3文件 ')


parser.add_argument('-o',
                help='必须给定，输出文件')


# 取出pasa中被 merge的 mRNA

args = parser.parse_args()

if not args.i or not args.o:
    parser.print_help()
    sys.exit()



infile = args.i

outfile= open(args.o,'w')

count_dic = {}


with open(infile) as fila:
	for i in fila:
		k = i.strip().split('\t')
		if len(k)>8:
			tmp_dic = {j.split('=')[0]:j.split('=')[1] for j in k[8].split(';') if j!=''}

			if len(tmp_dic['ID'].split('_'))>1:

				if k[2] == 'gene':
					#print(k)

					new_name = tmp_dic['ID'].split('_')[0]
					old_name = tmp_dic['ID'].split('_')[0]
					#M = ''

					if new_name not in count_dic:
						count_dic[new_name] = 1
					else:
						count_dic[new_name] += 1

					old_name = new_name
					new_name = new_name+'M'+str(count_dic[new_name])

				for j in tmp_dic:
					add_text = ''

					if tmp_dic[j].split('.')[-1][:4].upper()=='MRNA':
						add_text = tmp_dic[j].split('.')[-1]
					if tmp_dic[j].split('.')[-1][:3].upper()=='CDS':
						add_text = tmp_dic[j].split('.')[-1]
					if tmp_dic[j].split('.')[-1][:4].upper()=='EXON':
						add_text = tmp_dic[j].split('.')[-1]
					if tmp_dic[j].split('.')[-1][:3].upper()=='UTR':
						add_text = tmp_dic[j].split('.')[-1]
					tmp_dic[j] = (re.split(r'.mRNA[0-9]+',tmp_dic[j].split('%3B%20')[0].split('_')[0])[0].replace('MERGED%3A%20','').replace(old_name,new_name) +'.' +add_text).strip('.')

				k[8] = ';'.join([ j+'='+tmp_dic[j] for j in ['ID','Name', 'Parent'] if j in tmp_dic])

		outfile.write('\t'.join(k)+'\n')
			
				#break
outfile.close()